package com.github.ghsea.scheduler.util;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import lombok.Getter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.stereotype.Component;

/**
 * 所有的线程池都在这里统一管理
 *
 * @author guhai
 */
@Configuration
@Component
public class ThreadPoolConfig {

    private static Logger logger = LoggerFactory.getLogger(ThreadPoolConfig.class);

    private Integer cpuCount = Runtime.getRuntime().availableProcessors();

    @Bean("tinyTaskExecutor")
    public ExecutorService tinyTaskExecutor() {
        RejectedExecutionHandler rejectedExecutionHandler = new CountAndAbortPolicy();
        ThreadFactory threadFactory = newThreadFactory("tinyTaskExt-%s");
        return new ThreadPoolExecutor(cpuCount, cpuCount * 2, 10L, TimeUnit.MINUTES, new LinkedBlockingQueue<>(cpuCount * 4), threadFactory, rejectedExecutionHandler);
    }

    @Bean("scheduledExecutor")
    public ScheduledExecutorService scheduledExecutor() {
        RejectedExecutionHandler rejectedExecutionHandler = new CountAndAbortPolicy();
        ThreadFactory threadFactory = newThreadFactory("scheduledExt-%s");
        return new ScheduledThreadPoolExecutor(1, threadFactory, rejectedExecutionHandler);
    }

    /**
     * 专用请求中交兴路接口的线程池。不用使用该线程池做其他业务
     */
    @Bean("taskQueryExecutor")
    public ExecutorService taskQueryExecutor() {
        //如果线程处理不过来，使用CallerRunsPolicy起到反压效果，减缓消费速度
        RejectedExecutionHandler rejectedExecutionHandler = new CountAndCallerRunsPolicy();
        ThreadFactory threadFactory = newThreadFactory("QueryTaskExe-%s");
        return new ThreadPoolExecutor(cpuCount, cpuCount * 2,
                30L, TimeUnit.MINUTES, new LinkedBlockingQueue<>(cpuCount * 4), threadFactory, rejectedExecutionHandler);
    }

    private static ThreadFactory newThreadFactory(String nameFormat) {
        return new ThreadFactoryBuilder()
                .setNameFormat(nameFormat)
                .setDaemon(false)
                .setUncaughtExceptionHandler((t, e) -> logger.error("thread:" + t.getName() + " exception:", e)).build();
    }

    interface CountRejectionPolicy {
        long getRejectedTaskCount();
    }

    @Getter
    class CountAndAbortPolicy implements RejectedExecutionHandler, CountRejectionPolicy {
        private AtomicLong rejectedTasks = new AtomicLong(0);

        private RejectedExecutionHandler delegate = new ThreadPoolExecutor.AbortPolicy();

        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
            rejectedTasks.incrementAndGet();
            delegate.rejectedExecution(r, executor);
        }

        @Override
        public long getRejectedTaskCount() {
            return rejectedTasks.longValue();
        }
    }

    class CountAndCallerRunsPolicy implements RejectedExecutionHandler, CountRejectionPolicy {
        private AtomicLong rejectedTasks = new AtomicLong(0);

        private RejectedExecutionHandler delegate = new ThreadPoolExecutor.CallerRunsPolicy();

        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
            rejectedTasks.incrementAndGet();
            delegate.rejectedExecution(r, executor);
        }

        @Override
        public long getRejectedTaskCount() {
            return rejectedTasks.longValue();
        }
    }

}
